import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader


def test_img(model, test_data, args):
    """Evaluate model performance on test dataset."""
    model.eval()
    test_loss = 0.0
    correct = 0
    test_loader = DataLoader(test_data, batch_size=args.bs)

    with torch.no_grad():
        for data, target in test_loader:
            if args.gpu != -1:
                data, target = data.cuda(), target.cuda()

            output = model(data)
            test_loss += F.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    accuracy = 100.0 * correct / len(test_loader.dataset)

    if args.verbose:
        print(f'\nTest set: Average loss: {test_loss:.4f} '
              f'Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)\n')

    return accuracy, test_loss